/*
 * Copyright:  2018 smarabbit studio.
 *
 * Licensed under the Confluent Community License; you may not use this file except in
 * compliance with the License.  You may obtain a copy of the License at
 *
 * http://www.confluent.io/confluent-community-license
 *
 *  Unless required by applicable law or agreed to in writing, software distributed under
 *  the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 *  either express or implied.  See the License for the specific language governing
 *  permissions and limitations under the License.
 *
 *  @作   者： 黄开晖 (117227773@qq.com)
 *  @日   期:    2020/6/25 下午10:21
 *
 */

package com.massyframework.beanskin.graphql.servlet.provider;

import static java.util.Collections.singletonList;
import static java.util.stream.Collectors.toList;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.function.Supplier;
import java.util.stream.Stream;

import javax.websocket.HandshakeResponse;
import javax.websocket.server.HandshakeRequest;
import javax.websocket.server.ServerEndpointConfig;

import graphql.kickstart.execution.subscriptions.GraphQLSubscriptionMapper;
import graphql.kickstart.execution.subscriptions.SubscriptionProtocolFactory;
import graphql.kickstart.execution.subscriptions.apollo.ApolloSubscriptionConnectionListener;
import graphql.kickstart.servlet.GraphQLConfiguration;
import graphql.kickstart.servlet.GraphQLWebsocketServlet;
import graphql.kickstart.servlet.apollo.ApolloWebSocketSubscriptionProtocolFactory;
import graphql.kickstart.servlet.subscriptions.FallbackSubscriptionProtocolFactory;

/**
 * {@link GraphQLWSEndpointConfigurer}, Graphql WebSocket 端点配置
 */
public class GraphQLWSEndpointConfigurer extends ServerEndpointConfig.Configurator {

    private static final String HANDSHAKE_REQUEST_KEY = HandshakeRequest.class.getName();
    private static final String PROTOCOL_FACTORY_REQUEST_KEY = SubscriptionProtocolFactory.class.getName();

    private final List<SubscriptionProtocolFactory> subscriptionProtocolFactories;
    private final SubscriptionProtocolFactory fallbackSubscriptionProtocolFactory;
    private final List<String> allSubscriptionProtocols;
    
    private final Supplier<GraphQLConfiguration> configurationSupplier;

    public GraphQLWSEndpointConfigurer(Supplier<GraphQLConfiguration> configurationSupplier) {
    	this.configurationSupplier = Objects.requireNonNull(configurationSupplier, "\"configurationSupplier\" cannot be null.");
    	
    	GraphQLConfiguration configuration = this.configurationSupplier.get();
        List<ApolloSubscriptionConnectionListener> listeners = new ArrayList<>();
       
        subscriptionProtocolFactories = singletonList(new ApolloWebSocketSubscriptionProtocolFactory(
        		configuration.getObjectMapper(),
        		configuration.getInvocationInputFactory(),
        		configuration.getGraphQLInvoker(),
                listeners
        ));
        fallbackSubscriptionProtocolFactory = new FallbackSubscriptionProtocolFactory(
                new GraphQLSubscriptionMapper(configuration.getObjectMapper()),
                configuration.getInvocationInputFactory(),
                configuration.getGraphQLInvoker()
        );
        allSubscriptionProtocols = Stream
                .concat(subscriptionProtocolFactories.stream(), Stream.of(fallbackSubscriptionProtocolFactory))
                .map(SubscriptionProtocolFactory::getProtocol)
                .collect(toList());
    }

    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
        super.modifyHandshake(sec, request, response);

        sec.getUserProperties().put(HANDSHAKE_REQUEST_KEY, request);

        List<String> protocol = request.getHeaders().get(HandshakeRequest.SEC_WEBSOCKET_PROTOCOL);
        if (protocol == null) {
            protocol = Collections.emptyList();
        }

        SubscriptionProtocolFactory subscriptionProtocolFactory = getSubscriptionProtocolFactory(protocol);
        sec.getUserProperties().put(PROTOCOL_FACTORY_REQUEST_KEY, subscriptionProtocolFactory);

        if (request.getHeaders().get(HandshakeResponse.SEC_WEBSOCKET_ACCEPT) != null) {
            response.getHeaders().put(HandshakeResponse.SEC_WEBSOCKET_ACCEPT, allSubscriptionProtocols);
        }
        if (!protocol.isEmpty()) {
            response.getHeaders().put(HandshakeRequest.SEC_WEBSOCKET_PROTOCOL,
                    singletonList(subscriptionProtocolFactory.getProtocol()));
        }
    }

    @Override
    public <T> T getEndpointInstance(Class<T> endpointClass) throws InstantiationException {
        T result = null;
        if (endpointClass == GraphQLWebsocketServlet.class){
            result = endpointClass.cast(new GraphQLWebsocketServlet(this.configurationSupplier.get()));
        }
        return result;
    }

    private SubscriptionProtocolFactory getSubscriptionProtocolFactory(List<String> accept) {
        for (String protocol : accept) {
            for (SubscriptionProtocolFactory subscriptionProtocolFactory : subscriptionProtocolFactories) {
                if (subscriptionProtocolFactory.getProtocol().equals(protocol)) {
                    return subscriptionProtocolFactory;
                }
            }
        }

        return fallbackSubscriptionProtocolFactory;
    }
}